import numpy as np
import pandas as pd
import torch
import logging
from typing import List, Set, Dict, Tuple, Optional, Any

logger = logging.getLogger(__name__)


class DataProcessor:
    """
    指标数据处理类，负责将原始指标数据转换为模型训练所需的格式
    """

    def __init__(self, window_size: int = 6, max_gap: int = 60):
        """
        初始化数据处理器

        参数:
            window_size: 时间窗口大小，用于创建样本序列
            max_gap: 最大时间间隔（秒），用于判断时间序列是否连续
        """
        self.window_size = window_size
        self.max_gap = max_gap

    def create_metric_mapping(
        self, metrics_df: pd.DataFrame, useful_entities: Set[str]
    ) -> Dict[str, int]:
        """
        创建指标名称到特征索引的映射

        处理步骤:
        1. 过滤只保留有用实体的指标数据
        2. 提取唯一指标名称并按字母顺序排序
        3. 创建指标名称到索引的映射

        参数:
            metrics_df: 指标数据DataFrame
            useful_entities: 有用实体ID的集合

        返回:
            dict: 指标名称到特征索引的映射
        """
        if metrics_df.empty:
            return {}

        # 步骤1: 只处理有用实体的指标
        filtered_df = metrics_df[metrics_df["cmdb_id"].isin(useful_entities)]

        if filtered_df.empty:
            logger.warning("过滤后的指标数据为空，无法创建指标映射")
            return {}

        # 步骤2: 提取所有唯一的指标名称，并按照字符串排序确保顺序稳定
        all_metric_names = sorted(
            filtered_df["metric_name"].unique(), key=lambda x: str(x)
        )

        # 步骤3: 创建映射，索引严格按照排序后的顺序分配
        metric_mapping = {name: idx for idx, name in enumerate(all_metric_names)}
        logger.info(
            f"创建了 {len(metric_mapping)} 个指标到特征的映射，按照指标名称字母顺序排序"
        )
        return metric_mapping

    def ensure_datetime(self, metrics_df: pd.DataFrame) -> pd.DataFrame:
        """
        确保timestamp列为datetime类型

        参数:
            metrics_df: 指标数据DataFrame

        返回:
            pd.DataFrame: 处理后的DataFrame
        """
        df = metrics_df.copy()
        if not pd.api.types.is_datetime64_any_dtype(df["timestamp"]):
            df["timestamp"] = pd.to_datetime(df["timestamp"])
        return df

    def split_data_by_anomaly(self, experiment_data, window_size=None):
        """
        根据异常发生时间分割数据为训练集(异常前)和测试集(异常后)

        参数:
            experiment_data: ExperimentData对象
            window_size: 时间窗口大小（如果提供，将确保测试集包含足够的历史数据以维持窗口连续性）

        返回:
            tuple: (训练数据, 测试数据)
        """
        return experiment_data.split_data_by_anomaly_time(window_size)

    def convert_metrics_to_samples(
        self,
        metrics_df: pd.DataFrame,
        graph,
        cmdb_id_to_idx: Dict[str, int],
        metric_to_feature_idx: Dict[str, int],
        useful_entities: Set[str],
    ) -> List[Tuple]:
        """
        将指标数据转换为模型训练所需的样本格式

        处理步骤:
        1. 过滤数据，只保留有用实体的指标
        2. 按时间戳分组，为每个时间点构建特征矩阵
        3. 创建滑动窗口样本
        4. 转换为PyTorch张量格式

        参数:
            metrics_df: 指标数据DataFrame
            graph: 系统依赖图
            cmdb_id_to_idx: 实体ID到索引的映射
            metric_to_feature_idx: 指标名称到特征索引的映射
            useful_entities: 有用实体集合

        返回:
            List[Tuple]: 样本列表 [(timestamp, graph, features, targets), ...]
        """
        if metrics_df.empty:
            return []

        # 步骤1: 过滤数据，只保留有用实体的指标
        filtered_df = metrics_df[metrics_df["cmdb_id"].isin(useful_entities)]

        if filtered_df.empty:
            logger.warning("过滤后的指标数据为空，无法创建样本")
            return []

        # 步骤2: 按时间戳分组，为每个时间点构建特征矩阵
        timestamps = sorted(filtered_df["timestamp"].unique())

        time_features = {}
        channel_dim = len(metric_to_feature_idx)
        node_count = len(cmdb_id_to_idx)

        for ts in timestamps:
            df_ts = filtered_df[filtered_df["timestamp"] == ts]
            features = np.zeros((node_count, channel_dim))

            # 填充特征矩阵
            for _, row in df_ts.iterrows():
                cmdb_id = row["cmdb_id"]
                metric_name = row["metric_name"]
                value = row["value"]

                if cmdb_id not in cmdb_id_to_idx:
                    continue

                node_idx = cmdb_id_to_idx[cmdb_id]
                feature_idx = metric_to_feature_idx.get(metric_name)
                if feature_idx is not None:
                    features[node_idx, feature_idx] = value

            time_features[ts] = features

        # 步骤3: 创建滑动窗口样本
        samples = []

        for i in range(len(timestamps) - self.window_size):
            window_ts = timestamps[i : i + self.window_size]
            target_ts = timestamps[i + self.window_size - 1]

            # 检查时间窗口是否连续
            if any(
                (window_ts[j + 1] - window_ts[j]).total_seconds() > self.max_gap
                for j in range(len(window_ts) - 1)
            ):
                continue

            # 构建样本和目标
            sample_features = np.stack([time_features[ts] for ts in window_ts[:-1]])
            sample_targets = time_features[target_ts]

            # 步骤4: 转换为PyTorch张量
            sample_features_tensor = torch.tensor(sample_features, dtype=torch.float32)
            sample_targets_tensor = torch.tensor(sample_targets, dtype=torch.float32)

            # 添加样本
            samples.append(
                (
                    window_ts[-1].timestamp(),
                    graph,
                    sample_features_tensor,
                    sample_targets_tensor,
                )
            )

        logger.info(f"生成了 {len(samples)} 个样本")
        return samples
